/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*!
 * \file ffn.cc
 * \brief
 */
#include <graph/utils/type_utils.h>
#include <register/op_impl_registry.h>
#include "log/ops_log.h"

using namespace ge;
namespace ops {

ge::graphStatus InferShapeFFN(gert::InferShapeContext* context) {
  auto in_shape = context->GetInputShape(0);
  OPS_LOG_E_IF_NULL(context, in_shape, return ge::GRAPH_FAILED);
  auto out_shape = context->GetOutputShape(0);
  OPS_LOG_E_IF_NULL(context, out_shape, return ge::GRAPH_FAILED);
  *out_shape = *in_shape;
  return ge::GRAPH_SUCCESS;
}

ge::graphStatus InferDataTypeFFN(gert::InferDataTypeContext* context) {
  auto input_x_dtype = context->GetInputDataType(0);
  if (input_x_dtype == ge::DT_INT8) {
    auto attrs = context->GetAttrs();
    const int64_t* output_dtype = attrs->GetInt(2);
    if (output_dtype != NULL && *output_dtype == 1) {
      context->SetOutputDataType(0, ge::DT_BF16);
    } else {
      context->SetOutputDataType(0, ge::DT_FLOAT16);
    }
  } else {
    context->SetOutputDataType(0, input_x_dtype);
  }
  return GRAPH_SUCCESS;
}

IMPL_OP_INFERSHAPE(FFN).InferShape(InferShapeFFN)
    .InferDataType(InferDataTypeFFN);
}  // namespace ops
